package icu.zhhll.util.ratelimiter;

/**
 * 简单的令牌桶限流
 *
 * @author zh
 * @date 2024/6/19 16:14
 */
public class SimpleRateLimiter {

    /**
     * 桶的大小
     */
    private int limit;

    /**
     * 桶当前的token数量
     */
    private int tokens = 0;

    /**
     * 构造参数
     *
     * @param limit 最大令牌数量
     * @param speed 每秒增加的令牌数量
     */
    public SimpleRateLimiter(int limit, int speed) {
        //初始化桶的大小，且桶一开始是满的
        this.limit = limit;
        this.tokens = this.limit;

        //任务线程：每秒新增speed个令牌
        asyncTask(speed);
    }

    /**
     * 异步任务
     */
    public void asyncTask(Integer speed) {
        while (true) {
            try {
                Thread.sleep(1000L);

                int newTokens = tokens + speed;
                if (newTokens > limit) {
                    tokens = limit;
                    System.out.println("令牌桶满了！！！");
                } else {
                    tokens = newTokens;
                }
            } catch (InterruptedException e) {
                //输出到日志文件中

            }
        }
    }

    /**
     * 根据令牌数判断是否允许执行，需要加锁
     */
    public synchronized boolean execute() {
        if (tokens > 0) {
            tokens = tokens - 1;
            return true;
        }
        return false;
    }

}
